/*
 * Copyright 2010 Nicholas Alexander Cross
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.dataporter.io.jdbc;

import org.apache.log4j.Logger;
import org.dataporter.data.DataObject;
import org.dataporter.io.DataEntity;
import org.dataporter.io.DataModel;
import org.dataporter.io.DataReader;
import org.dataporter.io.EntitySource;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.*;

public class JDBCDataReader<Source extends EntitySource<? extends Connection>> implements DataReader<Source> {

    private static final Logger logger = Logger.getLogger( JDBCDataReader.class );

    private static final int BATCH_SIZE = 500;

    private final DataModel model;

    private Map<String, Long> highestPrimaryKeyForEntity;
    private EntitySource<? extends Connection> source;

    private Connection currentConnection;
    private String currentEntityName;
    private long lastBatchIndex;

    private Iterator<String> entityIterator;
    private Iterator<DataObject> fetchedObjects;

    public JDBCDataReader( DataModel model ) {
        this.model = model;
    }

    public void init( Source source ) throws Exception {
        this.source = source;
        highestPrimaryKeyForEntity = new HashMap<String, Long>();
        entityIterator = source.entityNames().iterator();
        lastBatchIndex = 0;
    }

    public DataObject read() {

        try
        {
            Connection connection = currentReader();

            if ( connection != null )
            {
                if ( fetchedObjects == null )
                {
                    fetchDataObjects( connection );
                }

                if ( fetchedObjects != null && fetchedObjects.hasNext() )
                {
                    return fetchedObjects.next();
                }
                else
                {
                    fetchDataObjects( connection );

                    if ( fetchedObjects.hasNext() )
                    {
                        return read();
                    }
                    else if ( entityIterator.hasNext() )
                    {
                        currentConnection = null;
                        fetchedObjects = null;
                        lastBatchIndex = 0;
                        return read();
                    }
                }
            }
        }
        catch ( SQLException e )
        {
            logger.error( "Failed to read from connection", e );
        }

        return null;
    }

    private Connection connectionForEntityName( String entityName ) {
        return source.sourceForEntity( entityName );
    }

    private Connection currentReader() throws SQLException {
        if ( currentConnection == null )
        {
            if ( entityIterator.hasNext() )
            {
                currentEntityName = entityIterator.next();
                currentConnection = connectionForEntityName( currentEntityName );
            }
        }

        return currentConnection;
    }

    private void fetchDataObjects( Connection connection ) throws SQLException {
        DataEntity entity = model.getEntity( currentEntityName );
        String keyName = entity.batchFetchKeyName();

        if ( keyName != null )
        {
            Long maximumBatchIndex = getMaximumBatchIndex( connection, keyName );
            if ( lastBatchIndex < maximumBatchIndex )
            {
                long nextBatchIndex = lastBatchIndex + BATCH_SIZE;
                fetchDataObjectsInRange( connection, entity, lastBatchIndex, nextBatchIndex > maximumBatchIndex ? maximumBatchIndex : nextBatchIndex );
                lastBatchIndex = nextBatchIndex + 1;
            }
        }
        else if ( fetchedObjects != null )
        {
            String query = new StringBuilder( "SELECT " ).append( columnsForEntity( entity ) ).append( " FROM " ).append( currentEntityName ).append( ";" ).toString();
            ResultSet results = connection.createStatement().executeQuery( query );
            fetchedObjects = dataObjectsFromResultSet( results );
        }
    }

    private void fetchDataObjectsInRange( Connection connection, DataEntity entity, long from, long to ) throws SQLException {
        String keyName = entity.batchFetchKeyName();

        String query = new StringBuilder( "SELECT " ).append( columnsForEntity( entity ) ).append( " FROM " ).append( currentEntityName )
                .append( " WHERE " ).append( keyName ).append( " >=" ).append( from )
                .append( " AND " ).append( keyName ).append( " <= " ).append( to ).append( ";" ).toString();

        ResultSet results = connection.createStatement().executeQuery( query );
        fetchedObjects = dataObjectsFromResultSet( results );
    }

    private Iterator<DataObject> dataObjectsFromResultSet( ResultSet results ) throws SQLException {
        List<DataObject> objects = new ArrayList<DataObject>();

        ResultSetMetaData metaData = results.getMetaData();
        int columnCount = metaData.getColumnCount();

        while (results.next()) {
            DataObject o = new DataObject( currentEntityName, model.getEntity( currentEntityName ).getPrimaryKeyNames() );

            for (int i=1; i<columnCount+1; i++) {
                String columnName = metaData.getColumnName( i );
                o.put( columnName, results.getObject( i ));
            }

            objects.add( o );
        }

        results.close();

        return objects.iterator();
    }

    private String columnsForEntity( DataEntity entity ) {
        String columnsToFetch = "*";

        List<String> columnNames = entity.getColumnNames();

        if ( columnNames != null )
        {
            StringBuilder builder = new StringBuilder();
            int i = 0;

            for ( String columnName : columnNames )
            {
                builder.append( columnName );
                i++;

                if ( i < columnNames.size() )
                {
                    builder.append( "," );
                }
            }
        }

        return columnsToFetch;
    }

    private long getMaximumBatchIndex( Connection connection, String keyName ) throws SQLException {
        Long highestPrimaryKey = highestPrimaryKeyForEntity.get( currentEntityName );

        if ( highestPrimaryKey == null )
        {
            String query = new StringBuilder( "SELECT MAX(" ).append( keyName ).append( ") FROM " ).append( currentEntityName ).append( ";" ).toString();
            ResultSet resultSet = connection.createStatement().executeQuery( query );
            resultSet.next();
            highestPrimaryKey = resultSet.getLong( 1 );
            highestPrimaryKeyForEntity.put( currentEntityName, highestPrimaryKey );
        }

        return highestPrimaryKey;
    }

    public void close() {
        Set<Connection> allConnections = new HashSet<Connection>();

        for ( String entityName : source.entityNames() )
        {
            Connection connection = source.sourceForEntity( entityName );
            allConnections.add( connection );
        }

        for ( Connection c : allConnections )
        {
            try
            {
                if ( !c.isClosed() )
                {
                    c.close();
                }
            }
            catch ( Exception e )
            {
                logger.error( "Failed to correctly close database connection", e );
            }
        }

        source = null;
        highestPrimaryKeyForEntity = null;
        fetchedObjects = null;
        entityIterator = null;
        lastBatchIndex = 0;
    }
}
